#include <jni.h>
#include <jni_md.h>
#include <string.h>
#include <string>
#include <iostream>
#include <queue>

#include "ffcodec.h"

//简单调用测试
extern "C"
JNIEXPORT jstring JNICALL
Java_com_wex_jni_JniFFmpeg_ffjniInfo(JNIEnv* env, jobject obj) {
    char *info = ffcodec_info();
    int len = strlen(info);
    jchar res[len];
    for (int i = 0; i < len; i++)
        res[i] = (jchar)info[i];
    return env->NewString(res, len);
}

// --- 编解码器本地句柄 ---

static const int Type_Encode = 0;
static const int Type_Decode = 1;

static const int Format_H264 = 0;
static const int Format_H265 = 1;

static const int Fifo_Max_Size = 10; //限制fifo最大缓冲数,防止爆炸

typedef struct {
    uint8_t *buff = NULL;
    int len = 0;
} CodecBuffer;

typedef struct {
    void *handle = NULL;
    CodecControlOhosType type;
    //yuv编码h26x时使用
    int mapSize = 0;
    FFCodec_Map map[3] = {};
    //编解码输出数据缓冲
    std::queue<CodecBuffer> fifo;
} CodecStruct;

// --- 编码接口 ---

extern "C"
JNIEXPORT jlong JNICALL
Java_com_wex_jni_JniFFmpeg_codecOpen(
        JNIEnv* env,
        jobject obj,
        jint type,
        jint format,
        jint width,
        jint height,
        jint fps)
{
    void *handle = NULL;

    //参数准备
    CodecControlOhosType codecType = type == Type_Encode ? CODEC_TYPE_ENCODE : CODEC_TYPE_DECODE;
    CodecControlOhosFormat inputFormat;
    CodecControlOhosFormat outputFormat;

    switch (format) {
        case Format_H264:
            if (codecType == CODEC_TYPE_ENCODE) {
                inputFormat = CODEC_FORMAT_YUV420P;
                outputFormat = CODEC_FORMAT_H264;
            } else {
                inputFormat = CODEC_FORMAT_H264;
                outputFormat = CODEC_FORMAT_YUV420P;
            }
            break;
        case Format_H265:
            if (codecType == CODEC_TYPE_ENCODE) {
                inputFormat = CODEC_FORMAT_YUV420P;
                outputFormat = CODEC_FORMAT_H265;
            } else {
                inputFormat = CODEC_FORMAT_H265;
                outputFormat = CODEC_FORMAT_YUV420P;
            }
            break;
        default:
            return -1;
    }

    //初始化编解码器
    handle = ffcodec_init(
            codecType,
            inputFormat,
            outputFormat,
            width,
            height,
            fps,
            20);

    if (!handle)
        return -2;

    //参数备份
    CodecStruct *cs = new CodecStruct();
    cs->handle = handle;
    cs->type = codecType;

    //yuv编码h26x时,传入yuv数据的结构体准备
    if (type == Type_Encode) {
        //传入yuv数据格式准备
        cs->mapSize = width * height * 3 / 2;

        cs->map[0].format = CODEC_FORMAT_YUV420P;
        cs->map[0].width = width;
        cs->map[0].height = height;
        cs->map[0].pb = 1;

        cs->map[1].format = CODEC_FORMAT_YUV420P;
        cs->map[1].width = width / 2;
        cs->map[1].height = height / 2;
        cs->map[1].pb = 1;

        cs->map[2].format = CODEC_FORMAT_YUV420P;
        cs->map[2].width = width / 2;
        cs->map[2].height = height / 2;
        cs->map[2].pb = 1;
    }
    else {
        ;
    }

    return (jlong)cs;
}

extern "C"
JNIEXPORT jint JNICALL
Java_com_wex_jni_JniFFmpeg_codecClose(JNIEnv* env, jobject obj, jlong fd)
{
    CodecStruct *cs = (CodecStruct*)fd;

    //内存回收
    if (fd > 0 && cs->handle != NULL) {
        //清空fifo
        ffcodec_deinit(cs->handle);
        while (!cs->fifo.empty()) {
            CodecBuffer buffer = cs->fifo.front();
            if (buffer.buff)
                delete[] buffer.buff;
            cs->fifo.pop();
        }
        //关闭编解码器
        delete cs;
    }

    return 0;
}

static CodecBuffer codecEncode(JNIEnv* env, CodecStruct *cs, jbyteArray data)
{
    CodecBuffer retBuffer = {};

    //格式转换
    jbyte *bytes = env->GetByteArrayElements(data, 0);
    int bytesSize = env->GetArrayLength(data);

    if (bytes && bytesSize > 0)
    {
        //转存yuv数据,长度达标的话直接使用内存
        if (bytesSize < cs->mapSize) {
            cs->map[0].map = new uint8_t[cs->mapSize];
            memcpy(cs->map[0].map, bytes, bytesSize);
        }
        else {
            cs->map[0].map = (uint8_t*)bytes;
        }

        //uv和y共用一段内存
        cs->map[1].map = cs->map[0].map + (cs->map[0].width * cs->map[0].height * cs->map[0].pb);
        cs->map[2].map = cs->map[1].map + (cs->map[1].width * cs->map[1].height * cs->map[1].pb);

        //开始编码
        uint8_t *frameBuff = NULL;
        uint32_t frameBuffSize = 0;
        if (ffcodec_encode(
            cs->handle,
            cs->map,
            &frameBuff,
            &frameBuffSize) == 0)
        {
            //生成返回数组
            if (frameBuff && frameBuffSize > 0)
            {
                retBuffer.len = frameBuffSize;
                retBuffer.buff = new uint8_t[frameBuffSize];
                memcpy(retBuffer.buff, frameBuff, frameBuffSize);
            }
        }

        //free
        if (cs->map[0].map != (uint8_t*)bytes)
            delete[] cs->map[0].map;
    }

    //free
    if (bytes)
        env->ReleaseByteArrayElements(data, bytes, 0);

    return retBuffer;
}

static CodecBuffer codecDecode(JNIEnv* env, CodecStruct *cs, jbyteArray data)
{
    CodecBuffer retBuffer = {};

    //格式转换
    jbyte *bytes = env->GetByteArrayElements(data, 0);
    int bytesSize = env->GetArrayLength(data);

    if (bytes && bytesSize > 0)
    {
        //开始解码
        if (ffcodec_decode(
            cs->handle,
            (uint8_t*)bytes,
            bytesSize,
            cs->map,
            NULL) == 0)
        {
            //计算yuv各包长度
            int ySize = cs->map[0].width * cs->map[0].height * cs->map[0].pb;
            int uSize = cs->map[1].width * cs->map[1].height * cs->map[1].pb;
            int vSize = cs->map[2].width * cs->map[2].height * cs->map[2].pb;
            //生成返回数组
            retBuffer.len = ySize + uSize + vSize;
            retBuffer.buff = new uint8_t[retBuffer.len];
            memcpy(&retBuffer.buff[0], cs->map[0].map, ySize);
            memcpy(&retBuffer.buff[ySize], cs->map[1].map, uSize);
            memcpy(&retBuffer.buff[ySize + vSize], cs->map[2].map, vSize);
        }
    }

    //free
    if (bytes)
        env->ReleaseByteArrayElements(data, bytes, 0);

    return retBuffer;
}

extern "C"
JNIEXPORT jint JNICALL
Java_com_wex_jni_JniFFmpeg_codecPush(JNIEnv* env, jobject obj, jlong fd, jbyteArray data)
{
    //非法fd
    CodecStruct *cs = (CodecStruct*)fd;
    if (fd < 1)
        return -1;

    //缓冲已满,请先codecGain()取走数据
    if (cs->fifo.size() >= Fifo_Max_Size)
        return -2;

    //编解码
    CodecBuffer buffer = {};
    if (cs->type == CODEC_TYPE_ENCODE)
        buffer = codecEncode(env, cs, data);
    else
        buffer = codecDecode(env, cs, data);

    if (buffer.len < 1)
        return -3;

    //把输出数据转存到fifo缓冲区,调用codecGain()可取走数据
    cs->fifo.push(buffer);
    return 0;
}

extern "C"
JNIEXPORT jbyteArray JNICALL
Java_com_wex_jni_JniFFmpeg_codecGain(JNIEnv* env, jobject obj, jlong fd)
{
    //非法fd
    CodecStruct *cs = (CodecStruct*)fd;
    if (fd < 1)
        return nullptr;

    //从缓冲区取数据
    if (cs->fifo.empty())
        return nullptr;

    CodecBuffer buffer = cs->fifo.front();
    cs->fifo.pop();
    if (buffer.len < 1)
        return nullptr;

    //合成java用的数据结构jbyteArray,对应java代码中的byte[]
    jbyteArray ret = env->NewByteArray(buffer.len);
    env->SetByteArrayRegion(ret, 0, buffer.len, reinterpret_cast<jbyte*>(buffer.buff));
    delete[] buffer.buff;

    return ret;
}
